import torch
import torch.nn as nn
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
import torchvision.datasets
from torch.utils.data import random_split
import torch.optim as optim

import pandas as pd 
import os 
import numpy as np 
import matplotlib.pyplot as plt

from skimage.transform import resize 
from skimage.io import imread  
from sklearn.model_selection import GridSearchCV 
from sklearn.model_selection import train_test_split 
from sklearn.metrics import accuracy_score 
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix
from sklearn.metrics import ConfusionMatrixDisplay

import seaborn as sns # Beautify CM

from torchvision.models import ResNet50_Weights, VGG16_Weights, EfficientNet_B0_Weights

# Transformations for the images
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize to 224x224
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # Normalize using ImageNet stats
])

# Load dataset
dataset_path = "C:/Users/rishi/Desktop/JHU/Critical Infrastructure Protection/Major Project/Data Sets/Sonar/TrainSetMotionBlur"
dataset = torchvision.datasets.ImageFolder(root=dataset_path, transform=transform)

# Train-test split
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

# DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# Classes
class_names = dataset.classes
num_classes = len(class_names)

print(class_names, num_classes)

['BigAnimals', 'Mines', 'Pipes', 'Rockets', 'Vehicles'] 5

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Running on: ", device)

Running on:  cuda

# Set seeds for reproducibility
torch.manual_seed(42)  # For PyTorch
np.random.seed(42)     # For NumPy

def setup_model_with_dropout(model_type, num_classes, dropout_rate=0.0):
    if model_type == 'resnet':
        model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        model.fc = nn.Sequential(
            nn.Dropout(dropout_rate),  # Add dropout before final FC
            nn.Linear(model.fc.in_features, num_classes)
        )

    elif model_type == 'vgg':
        model = models.vgg16(weights=models.VGG16_Weights.DEFAULT)
        model.classifier[6] = nn.Sequential(
            nn.Dropout(dropout_rate),  # Add dropout before final FC
            nn.Linear(model.classifier[6].in_features, num_classes)
        )

    elif model_type == 'efficientnet':
        model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)
        model.classifier = nn.Sequential(
            nn.Dropout(dropout_rate),  # Add dropout before final FC
            nn.Linear(model.classifier[1].in_features, num_classes)
        )

    model = model.to(device)
    return model

def setup_optimizer(model, optimizer_name, learning_rate, weight_decay):
    if optimizer_name == 'adam':
        return optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    elif optimizer_name == 'sgd':
        return optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay)
    elif optimizer_name == 'rmsprop':
        return optim.RMSprop(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    else:
        raise ValueError(f"Unsupported optimizer: {optimizer_name}")

def train_and_evaluate(model, optimizer, criterion, train_loader, test_loader, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(train_loader)}")

    model.eval()
    y_true, y_pred = [], []
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predictions = torch.max(outputs, 1)
            y_true.extend(labels.cpu().numpy())
            y_pred.extend(predictions.cpu().numpy())

    accuracy = accuracy_score(y_true, y_pred)
    cm = confusion_matrix(y_true, y_pred)
    report = classification_report(y_true, y_pred, target_names=class_names, zero_division = 1)

    print(f"Accuracy: {accuracy * 100:.2f}%")
    print("Classification Report:\n", report)

    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.show()

    return accuracy, report, cm

def hyperparameter_search(model_types, param_grid, train_dataset, test_dataset, num_classes):
    best_model, best_metrics, best_params = None, {'accuracy': 0}, None

    for model_type in model_types:
        for lr in param_grid['learning_rate']:
            for batch_size in param_grid['batch_size']:
                for opt in param_grid['optimizer']:
                    for wd in param_grid['weight_decay']:
                        for dr in param_grid['dropout_rate']:
                            print(f"Training {model_type} with lr={lr}, batch_size={batch_size}, "
                                  f"optimizer={opt}, weight_decay={wd}, dropout_rate={dr}")

                            # Create DataLoaders
                            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
                            test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

                            # Set up the model with the specified dropout rate
                            model = setup_model_with_dropout(model_type, num_classes, dr)

                            # Set up the optimizer
                            optimizer = setup_optimizer(model, opt, lr, wd)

                            # Define the loss function
                            criterion = nn.CrossEntropyLoss()

                            # Train and evaluate
                            accuracy, _, _ = train_and_evaluate(
                                model, optimizer, criterion, train_loader, test_loader
                            )

                            # Track the best model
                            if accuracy > best_metrics['accuracy']:
                                best_model = model
                                best_metrics['accuracy'] = accuracy
                                best_params = {
                                    'model': model_type,
                                    'learning_rate': lr,
                                    'batch_size': batch_size,
                                    'optimizer': opt,
                                    'weight_decay': wd,
                                    'dropout_rate': dr
                                }

    print(f"Best Model: {best_params['model']} | Accuracy: {best_metrics['accuracy']:.2f}% | Params: {best_params}")
    return best_model, best_metrics, best_params

param_grid = {
    'learning_rate': [0.001, 0.0005, 0.0001],
    'batch_size': [16, 32, 64],
    'optimizer': ['adam', 'sgd', 'rmsprop'],
    'weight_decay': [0.0, 1e-4, 1e-3],
    'dropout_rate': [0.0, 0.3, 0.5]
}

model_types = ['resnet', 'vgg', 'efficientnet']

best_model, best_metrics, best_params = hyperparameter_search(
    model_types, param_grid, train_dataset, test_dataset, num_classes
)

# Save the best model
torch.save(best_model.state_dict(), f"best_model_{best_params['model']}.pth")
print("Best model saved successfully!")

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.0, dropout_rate=0.0
Epoch 1/10, Loss: 1.4226167268223233
Epoch 2/10, Loss: 1.064807312356101
Epoch 3/10, Loss: 0.8907065325313144
Epoch 4/10, Loss: 0.8274083733558655
Epoch 5/10, Loss: 0.6042677437265714
Epoch 6/10, Loss: 0.6789631719390551
Epoch 7/10, Loss: 0.5500290770497587
Epoch 8/10, Loss: 0.4349922442601787
Epoch 9/10, Loss: 0.4354906040761206
Epoch 10/10, Loss: 0.32993256486952305
Accuracy: 70.00%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.75      0.55      0.63        11
       Mines       0.75      0.63      0.69        19
       Pipes       0.88      0.82      0.85        17
     Rockets       0.45      0.91      0.61        11
    Vehicles       0.88      0.58      0.70        12

    accuracy                           0.70        70
   macro avg       0.74      0.70      0.69        70
weighted avg       0.76      0.70      0.71        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.0, dropout_rate=0.3
Epoch 1/10, Loss: 1.162516188290384
Epoch 2/10, Loss: 0.9404450787438287
Epoch 3/10, Loss: 0.5829256524642309
Epoch 4/10, Loss: 0.738079573545191
Epoch 5/10, Loss: 0.4737313182817565
Epoch 6/10, Loss: 0.2142927392075459
Epoch 7/10, Loss: 0.15698119439184666
Epoch 8/10, Loss: 0.30264253324518603
Epoch 9/10, Loss: 0.3030023028453191
Epoch 10/10, Loss: 0.276163334854775
Accuracy: 74.29%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.69      0.82      0.75        11
       Mines       1.00      0.42      0.59        19
       Pipes       0.61      1.00      0.76        17
     Rockets       0.80      0.73      0.76        11
    Vehicles       0.91      0.83      0.87        12

    accuracy                           0.74        70
   macro avg       0.80      0.76      0.75        70
weighted avg       0.81      0.74      0.73        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.0, dropout_rate=0.5
Epoch 1/10, Loss: 1.3983533051278856
Epoch 2/10, Loss: 0.7916058633062575
Epoch 3/10, Loss: 0.8933926026026408
Epoch 4/10, Loss: 0.6838194727897644
Epoch 5/10, Loss: 0.6425449649492899
Epoch 6/10, Loss: 0.5057451609108183
Epoch 7/10, Loss: 0.48947153488794964
Epoch 8/10, Loss: 0.2554271484000815
Epoch 9/10, Loss: 0.2746824228929149
Epoch 10/10, Loss: 0.31304313087215024
Accuracy: 78.57%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.56      0.91      0.69        11
       Mines       0.88      0.79      0.83        19
       Pipes       0.94      0.88      0.91        17
     Rockets       0.64      0.64      0.64        11
    Vehicles       1.00      0.67      0.80        12

    accuracy                           0.79        70
   macro avg       0.80      0.78      0.77        70
weighted avg       0.83      0.79      0.79        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.0001, dropout_rate=0.0
Epoch 1/10, Loss: 1.193955808877945
Epoch 2/10, Loss: 0.8686565442217721
Epoch 3/10, Loss: 0.7154402765962813
Epoch 4/10, Loss: 0.5626552287075255
Epoch 5/10, Loss: 0.4301767258180512
Epoch 6/10, Loss: 0.341320570351349
Epoch 7/10, Loss: 0.3723193659550614
Epoch 8/10, Loss: 0.2415904640737507
Epoch 9/10, Loss: 0.13053373568173912
Epoch 10/10, Loss: 0.2285459830891341
Accuracy: 82.86%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.88      0.64      0.74        11
       Mines       0.76      0.84      0.80        19
       Pipes       1.00      0.94      0.97        17
     Rockets       0.75      0.82      0.78        11
    Vehicles       0.77      0.83      0.80        12

    accuracy                           0.83        70
   macro avg       0.83      0.81      0.82        70
weighted avg       0.84      0.83      0.83        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.0001, dropout_rate=0.3
Epoch 1/10, Loss: 1.3468499349223242
Epoch 2/10, Loss: 0.9414312210347917
Epoch 3/10, Loss: 0.7709920803705851
Epoch 4/10, Loss: 0.5886638263861338
Epoch 5/10, Loss: 0.3070489809744888
Epoch 6/10, Loss: 0.4753173639376958
Epoch 7/10, Loss: 0.3401246445460452
Epoch 8/10, Loss: 0.18201635653773943
Epoch 9/10, Loss: 0.19121312893306217
Epoch 10/10, Loss: 0.2542307372722361
Accuracy: 65.71%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       1.00      0.18      0.31        11
       Mines       0.50      0.89      0.64        19
       Pipes       1.00      0.82      0.90        17
     Rockets       0.62      0.45      0.53        11
    Vehicles       0.67      0.67      0.67        12

    accuracy                           0.66        70
   macro avg       0.76      0.60      0.61        70
weighted avg       0.75      0.66      0.64        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.0001, dropout_rate=0.5
Epoch 1/10, Loss: 1.3025224374400244
Epoch 2/10, Loss: 0.8990547673569785
Epoch 3/10, Loss: 0.7085604270299276
Epoch 4/10, Loss: 0.5290212598111894
Epoch 5/10, Loss: 0.4095746684405539
Epoch 6/10, Loss: 0.5472410544753075
Epoch 7/10, Loss: 0.2965755129439963
Epoch 8/10, Loss: 0.29717597189462847
Epoch 9/10, Loss: 0.2295161440140671
Epoch 10/10, Loss: 0.212235393623511
Accuracy: 77.14%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.89      0.73      0.80        11
       Mines       0.63      0.89      0.74        19
       Pipes       1.00      0.82      0.90        17
     Rockets       0.67      0.55      0.60        11
    Vehicles       0.82      0.75      0.78        12

    accuracy                           0.77        70
   macro avg       0.80      0.75      0.76        70
weighted avg       0.80      0.77      0.77        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.001, dropout_rate=0.0
Epoch 1/10, Loss: 1.2990308105945587
Epoch 2/10, Loss: 0.9530814819865756
Epoch 3/10, Loss: 0.5953774419095781
Epoch 4/10, Loss: 0.48136866754955715
Epoch 5/10, Loss: 0.4698384954697556
Epoch 6/10, Loss: 0.41774962345759076
Epoch 7/10, Loss: 0.43096932851605946
Epoch 8/10, Loss: 0.34274400853448445
Epoch 9/10, Loss: 0.3461839560833242
Epoch 10/10, Loss: 0.2830766240755717
Accuracy: 71.43%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.45      0.91      0.61        11
       Mines       0.83      0.53      0.65        19
       Pipes       1.00      0.76      0.87        17
     Rockets       0.80      0.73      0.76        11
    Vehicles       0.69      0.75      0.72        12

    accuracy                           0.71        70
   macro avg       0.76      0.74      0.72        70
weighted avg       0.78      0.71      0.72        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.001, dropout_rate=0.3
Epoch 1/10, Loss: 1.3330308828088973
Epoch 2/10, Loss: 1.0557886792553797
Epoch 3/10, Loss: 0.9798804554674361
Epoch 4/10, Loss: 0.6880867994493909
Epoch 5/10, Loss: 0.5080874214569727
Epoch 6/10, Loss: 0.5979864440030522
Epoch 7/10, Loss: 0.34178847074508667
Epoch 8/10, Loss: 0.31436599377128815
Epoch 9/10, Loss: 0.5198746778898768
Epoch 10/10, Loss: 0.393432407743401
Accuracy: 71.43%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.83      0.45      0.59        11
       Mines       0.65      0.89      0.76        19
       Pipes       0.74      0.82      0.78        17
     Rockets       0.83      0.45      0.59        11
    Vehicles       0.69      0.75      0.72        12

    accuracy                           0.71        70
   macro avg       0.75      0.68      0.69        70
weighted avg       0.74      0.71      0.70        70

Training resnet with lr=0.001, batch_size=16, optimizer=adam, weight_decay=0.001, dropout_rate=0.5
Epoch 1/10, Loss: 1.3946256306436327
Epoch 2/10, Loss: 0.9117740227116479
Epoch 3/10, Loss: 0.9644500994020038
Epoch 4/10, Loss: 0.6329642915063434
Epoch 5/10, Loss: 0.5945399800936381
Epoch 6/10, Loss: 0.5637176624602742
Epoch 7/10, Loss: 0.48064622200197643
Epoch 8/10, Loss: 0.5740980307261149
Epoch 9/10, Loss: 0.5354663423366017
Epoch 10/10, Loss: 0.39513255490197075
Accuracy: 70.00%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.67      0.55      0.60        11
       Mines       0.59      1.00      0.75        19
       Pipes       0.93      0.82      0.88        17
     Rockets       0.67      0.73      0.70        11
    Vehicles       1.00      0.17      0.29        12

    accuracy                           0.70        70
   macro avg       0.77      0.65      0.64        70
weighted avg       0.77      0.70      0.67        70

Training resnet with lr=0.001, batch_size=16, optimizer=sgd, weight_decay=0.0, dropout_rate=0.0
Epoch 1/10, Loss: 1.6025307443406847
Epoch 2/10, Loss: 1.5405692325698004
Epoch 3/10, Loss: 1.4435244997342427
Epoch 4/10, Loss: 1.3230764932102628
Epoch 5/10, Loss: 1.1616113119655185
Epoch 6/10, Loss: 0.9617549545235105
Epoch 7/10, Loss: 0.7858256134721968
Epoch 8/10, Loss: 0.6490684035751555
Epoch 9/10, Loss: 0.48043371323082185
Epoch 10/10, Loss: 0.3300580067767037
Accuracy: 81.43%
Classification Report:
               precision    recall  f1-score   support

  BigAnimals       0.83      0.91      0.87        11
       Mines       0.85      0.89      0.87        19
       Pipes       1.00      0.76      0.87        17
     Rockets       0.56      0.82      0.67        11
    Vehicles       0.89      0.67      0.76        12

    accuracy                           0.81        70
   macro avg       0.83      0.81      0.81        70
weighted avg       0.85      0.81      0.82        70

# The above code block ran for 81764 seconds = 23 hours

torch.save(best_model, f"best_model_{best_params['model']}_complete.pth")
print("Best model saved successfully!")

Best model saved successfully!

resnet_model = models.resnet50(weights=ResNet50_Weights.DEFAULT)
resnet_model.fc = nn.Linear(resnet_model.fc.in_features, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(resnet_model.parameters(), lr=0.001)

print("Training and Evaluating ResNet:")
resnet_accuracy, resnet_report, resnet_cm = train_and_evaluate(
    resnet_model, optimizer, criterion, train_loader, test_loader, num_epochs=10
)

vgg_model = models.vgg16(weights=VGG16_Weights.DEFAULT)
vgg_model.classifier[6] = nn.Linear(vgg_model.classifier[6].in_features, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(vgg_model.parameters(), lr=0.001)

print("Training and Evaluating VGG:")
vgg_accuracy, vgg_report, vgg_cm = train_and_evaluate(
    vgg_model, optimizer, criterion, train_loader, test_loader, num_epochs=10
)

efficientnet_model = models.efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT)
efficientnet_model.classifier[1] = nn.Linear(efficientnet_model.classifier[1].in_features, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(efficientnet_model.parameters(), lr=0.001)

print("Training and Evaluating EfficientNet:")
efficientnet_accuracy, efficientnet_report, efficientnet_cm = train_and_evaluate(
    efficientnet_model, optimizer, criterion, train_loader, test_loader, num_epochs=10
)

print("Summary of Results:")
print(f"ResNet Accuracy: {resnet_accuracy * 100:.2f}%")
print(f"VGG Accuracy: {vgg_accuracy * 100:.2f}%")
print(f"EfficientNet Accuracy: {efficientnet_accuracy * 100:.2f}%")

Libraries¶

Preprocessing¶

Model Setup¶

Optimizer Setup¶

Training and Evaluation Function¶

Hyperparameter Search¶

Defining Hyperparameter Grid and Running the Models¶

When you save the entire model, you are saving:¶

-->The model architecture: The complete definition of the neural network, including its layers and structure.¶

-->The model parameters (weights and biases): The learned values from training.¶

The state dictionary only saves:¶

-->The model's parameters: This includes weights, biases, and any other trainable tensors.¶

Initial Individual Models -- Ignore¶

ResNet¶

VGG¶

EfficientNet¶

Print Results for All Models¶